package cn.iocoder.yudao.module.resource.api;

import cn.hutool.core.bean.BeanUtil;
import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.crypto.digest.DigestUtil;
import cn.iocoder.yudao.framework.common.util.json.JsonUtils;
import cn.iocoder.yudao.module.infra.embedding.EmbeddingServiceFactory;
import cn.iocoder.yudao.module.infra.vector.VectorServiceFactory;
import cn.iocoder.yudao.module.resource.api.enums.EnumAipError;
import cn.iocoder.yudao.module.resource.api.enums.MessageContentTypeEnum;
import cn.iocoder.yudao.module.resource.api.enums.MessageSourceTypeEnum;
import cn.iocoder.yudao.module.resource.api.enums.MessageVectorColumnEnum;
import cn.iocoder.yudao.module.resource.api.model.message.AudioDataDTO;
import cn.iocoder.yudao.module.resource.api.model.messagecollection.MessageCollectionDTO;
import cn.iocoder.yudao.module.resource.api.model.messagecollection.MessageCollectionResDTO;
import cn.iocoder.yudao.module.resource.api.model.messagecollection.VectorDoc;
import cn.iocoder.yudao.module.resource.dal.dataobject.messagecollection.MessageCollectionDO;
import cn.iocoder.yudao.module.resource.dal.mysql.messagecollection.MessageCollectionMapper;
import cn.iocoder.yudao.module.resource.dal.redis.ResourceLockRedisDAO;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
import org.springframework.stereotype.Service;

import java.util.*;
import java.util.stream.Collectors;

import static cn.iocoder.yudao.framework.common.exception.util.ServiceExceptionUtil.exception;

@Slf4j
@Service
public class MessageCollectionApiImpl implements IMessageCollectionApi {
    @Resource
    private EmbeddingServiceFactory embeddingServiceFactory;
    @Resource
    private VectorServiceFactory vectorServiceFactory;
    @Resource
    private MessageCollectionMapper messageCollectionMapper;
    @Resource
    private ResourceLockRedisDAO resourceLockRedisDAO;

    @NotNull
    @Override
    public Map<String, List<? extends Number>> getVectorValues(String userText, String replyText) {
        final Map<String, List<? extends Number>> vectorValues = new HashMap<>();

        if (StrUtil.isNotEmpty(userText)) {
            final List<? extends Number> vUserText = embeddingServiceFactory.getService().create(userText);
            if (CollUtil.isEmpty(vUserText)) {
                throw exception(EnumAipError.MESSAGE_COLLECTION_DEFAULT_ERROR, "计算userText向量失败");
            }
            vectorValues.put(MessageVectorColumnEnum.V_USER_TEXT.getValue(), vUserText);
        }

        if (StrUtil.isNotEmpty(replyText)) {
            final List<? extends Number> vReplyText = embeddingServiceFactory.getService().create(replyText);
            if (CollUtil.isEmpty(vReplyText)) {
                throw exception(EnumAipError.MESSAGE_COLLECTION_DEFAULT_ERROR, "计算replyText向量失败");
            }
            vectorValues.put(MessageVectorColumnEnum.V_REPLY_TEXT.getValue(), vReplyText);
        }
        return vectorValues;
    }

    @Override
    public MessageCollectionResDTO queryMessage(Float similarity, int topK,
                                                String userText, String replyText) {
        if (similarity == null) {
            throw exception(EnumAipError.MESSAGE_COLLECTION_DEFAULT_ERROR, "similarity不能为空");
        }

        if (StrUtil.isEmpty(userText) && StrUtil.isEmpty(replyText)) {
            throw exception(EnumAipError.MESSAGE_COLLECTION_DEFAULT_ERROR, "userText/replyText不能同时为空");
        }

        // 计算向量
        final Map<String, List<? extends Number>> vectorValues = getVectorValues(userText, replyText);

        final List<MessageCollectionDTO> messages;
        if (StrUtil.isNotEmpty(userText) && StrUtil.isNotEmpty(replyText)) {
            // 先查询replyText
            final List<MessageCollectionDTO> replyTextMessages = queryDoc(similarity, topK,
                    Map.of(MessageVectorColumnEnum.V_REPLY_TEXT.getValue(), vectorValues.get(MessageVectorColumnEnum.V_REPLY_TEXT.getValue())), null);
            log.info("先查询replyText向量, userText = {}, replyText = {}, replyTextMessages = {}",
                    userText, replyText, JsonUtils.toJsonString(replyTextMessages));
            if (CollUtil.isEmpty(replyTextMessages)) {
                return new MessageCollectionResDTO(new ArrayList<>(), vectorValues);
            }

            // 再查询userText
            final List<String> queryMessageIds = replyTextMessages.stream().map(v -> "'" + v.getMessageId() + "'").collect(Collectors.toList());
            messages = queryDoc(similarity, topK,
                    Map.of(MessageVectorColumnEnum.V_USER_TEXT.getValue(), vectorValues.get(MessageVectorColumnEnum.V_USER_TEXT.getValue())),
                    StrUtil.format("{} in ({})", MessageVectorColumnEnum.MESSAGE_ID.getValue(), String.join(",", queryMessageIds)));
        } else {
            messages = queryDoc(similarity, topK, vectorValues, null);
        }

        log.info("查询向量, userText = {}, replyText = {}, messages = {}",
                userText, replyText, JsonUtils.toJsonString(messages));
        return new MessageCollectionResDTO(messages, vectorValues);
    }

    @NotNull
    private List<MessageCollectionDTO> queryDoc(Float similarity, int topK,
                                                Map<String, List<? extends Number>> vectorValues,
                                                String filter) {
        final List<VectorDoc> docs = vectorServiceFactory.getService().queryDoc(
                null, vectorValues, null, filter, topK, similarity,
                Arrays.asList(MessageVectorColumnEnum.MESSAGE_ID.getValue(),
                        MessageVectorColumnEnum.USER_TEXT.getValue(),
                        MessageVectorColumnEnum.REPLY_TEXT.getValue(),
                        MessageVectorColumnEnum.AUDIO_DATA.getValue()
                )
        );

        final List<MessageCollectionDTO> messages = new ArrayList<>();
        if (CollUtil.isNotEmpty(docs)) {
            for (VectorDoc doc : docs) {
                final String audioDataStr = (String) doc.getFields().get(MessageVectorColumnEnum.AUDIO_DATA.getValue());
                final List<AudioDataDTO> audioData = JsonUtils.parseArray(audioDataStr, AudioDataDTO.class);

                final MessageCollectionDTO message = new MessageCollectionDTO();
                BeanUtil.fillBeanWithMap(doc.getFields(), message, true);
                message.setAudioData(audioData);
                message.setSimilarity(doc.getSimilarity());
                messages.add(message);
            }
        }
        return messages;
    }

    @Override
    public void insertMessage(Integer toneId, MessageSourceTypeEnum messageSendType,
                              String messageId, List<AudioDataDTO> audioDataList,
                              String userText, String replyText,
                              Map<String, List<? extends Number>> vectorValues) {
        final Map<String, Object> fields = createMessageMap(messageSendType, messageId, audioDataList, userText, replyText, null);

        vectorServiceFactory.getService().insertDoc(messageId, fields, vectorValues);
        log.info("插入数据到向量数据库, messageId = {}, userText = {}, replyText = {}", messageId, userText, replyText);

        insertOrUpdateToDB(messageSendType, messageId, audioDataList, userText, replyText);
    }

    @Override
    public void updateMessage(Integer toneId, MessageSourceTypeEnum messageSendType,
                              String messageId, List<AudioDataDTO> audioDataList,
                              String userText, String replyText,
                              Map<String, List<? extends Number>> vectorValues) {
        resourceLockRedisDAO.lockMessage(messageId, () -> {
            final MessageCollectionDO message = messageCollectionMapper.selectById(messageId);
            if (message == null) {
                log.warn("更新数据到向量数据库-mysql中数据不存在, messageId = {}, userText = {}, replyText = {}", messageId, userText, replyText);
                return null;
            }

            final List<AudioDataDTO> existedAudioDataList = message.getAudioData();
            final Map<String, Object> fields = createMessageMap(messageSendType, messageId, audioDataList, userText, replyText, existedAudioDataList);

            vectorServiceFactory.getService().updateDoc(messageId, fields, vectorValues);
            log.info("更新数据到向量数据库, messageId = {}, userText = {}, replyText = {}", messageId, userText, replyText);

            insertOrUpdateToDB(messageSendType, messageId, audioDataList, userText, replyText);
            return null;
        });
    }

    private void insertOrUpdateToDB(MessageSourceTypeEnum messageSendType, String messageId,
                                    List<AudioDataDTO> audioData,
                                    String userText, String replyText) {
        final MessageCollectionDO messageCollection = MessageCollectionDO.builder()
                .messageId(messageId)
                .senderType(messageSendType.getValue())
                .userText(userText)
                .userTextHash(DigestUtil.md5Hex(userText))
                .replyText(replyText)
                .replyTextHash(DigestUtil.md5Hex(replyText))
                .contentType(MessageContentTypeEnum.TEXT_OR_AUDIO.getValue())
                .audioData(audioData != null ? audioData : new ArrayList<>())
                .build();
        messageCollectionMapper.insertOrUpdate(messageCollection);
        log.info("写入数据到消息仓库, messageId = {}, userText = {}, replyText = {}", messageId, userText, replyText);
    }

    @NotNull
    private static Map<String, Object> createMessageMap(MessageSourceTypeEnum messageSendType,
                                                        String messageId, List<AudioDataDTO> audioData,
                                                        String userText, String replyText,
                                                        List<AudioDataDTO> existedAudioDataList) {
        final List<AudioDataDTO> finalAudioList = CollUtil.isNotEmpty(audioData) ? new ArrayList<>(audioData) : new ArrayList<>();
        if (CollUtil.isNotEmpty(existedAudioDataList)) {
            final List<String> ids = finalAudioList.stream().map(v -> v.getId()).collect(Collectors.toList());
            for (AudioDataDTO dataDTO : existedAudioDataList) {
                if (ids.contains(dataDTO.getId())) {
                    continue;
                }

                finalAudioList.add(dataDTO);
                ids.add(dataDTO.getId());
            }
        }

        final Map<String, Object> fields = new HashMap<>();
        fields.put(MessageVectorColumnEnum.MESSAGE_ID.getValue(), messageId);
        fields.put(MessageVectorColumnEnum.SENDER_TYPE.getValue(), messageSendType.getValue());
        fields.put(MessageVectorColumnEnum.CONTENT_TYPE.getValue(), MessageContentTypeEnum.TEXT_OR_AUDIO.getValue());
        fields.put(MessageVectorColumnEnum.AUDIO_DATA.getValue(), CollUtil.isNotEmpty(finalAudioList) ? JsonUtils.toJsonString(finalAudioList) : "[]");
        fields.put(MessageVectorColumnEnum.ROLE_ID.getValue(), ""); // delete
        fields.put(MessageVectorColumnEnum.CREATE_TIME.getValue(), System.currentTimeMillis());
        fields.put(MessageVectorColumnEnum.USER_TEXT.getValue(), userText);
        fields.put(MessageVectorColumnEnum.REPLY_TEXT.getValue(), replyText);
        fields.put(MessageVectorColumnEnum.EXTENSION_STR1.getValue(), "");
        fields.put(MessageVectorColumnEnum.EXTENSION_STR2.getValue(), "");
        fields.put(MessageVectorColumnEnum.EXTENSION_STR3.getValue(), "");
        fields.put(MessageVectorColumnEnum.EXTENSION_INT1.getValue(), 0);
        fields.put(MessageVectorColumnEnum.EXTENSION_INT2.getValue(), 0);
        fields.put(MessageVectorColumnEnum.EXTENSION_INT3.getValue(), 0);
        fields.put(MessageVectorColumnEnum.EXTENSION_LONG1.getValue(), 0L);
        fields.put(MessageVectorColumnEnum.EXTENSION_LONG2.getValue(), 0L);
        fields.put(MessageVectorColumnEnum.EXTENSION_LONG3.getValue(), 0L);
        fields.put(MessageVectorColumnEnum.EXTENSION_FLOAT1.getValue(), 0f);
        fields.put(MessageVectorColumnEnum.EXTENSION_FLOAT2.getValue(), 0f);
        fields.put(MessageVectorColumnEnum.EXTENSION_FLOAT3.getValue(), 0f);
        return fields;
    }

    @Override
    public void deleteMessage(String messageId) {
        vectorServiceFactory.getService().deleteDoc(messageId);
        log.info("从向量数据库删除数据, messageId = {}", messageId);

        messageCollectionMapper.deleteById(messageId);
        log.info("从消息仓库删除数据, messageId = {}", messageId);
    }

    public static void main(String[] args) {
        List<String> messageIds = Arrays.asList("1", "2", "3");
        messageIds = messageIds.stream().map(v -> "'" + v + "'").collect(Collectors.toList());
        String str = StrUtil.format("{} in ({})", MessageVectorColumnEnum.MESSAGE_ID.getValue(), String.join(",", messageIds));
        System.out.println(str);
    }
}
